/*letkf4d.c
Implement the 3D-LETKF algorithm in parallel

Author: Shiming Yang
Date: Nov. 2009
*/
#include "letkf_mod.h"

void letkf_local(int num_ensemble, int radius, int numpar, int numline, int numobs, double *xb_bar, float *xb_tlat, float *xb_tlon, float *xb_zlvl, double *yb_bar, float *yb_clat, float *yb_clon, double *obs, float *obs_lat, float *obs_lon, double **en_bar, double **en_mass, double **en_con)

{
	int i = 0; int k = 0;
	int len_L = 0;
	double rho = 1.15;
	double valid_dist = 25.0;
	double sqrtcon = sqrt(num_ensemble - 1);
	int center = 0; int lenr = 0;
	int *len_r = &lenr;
	int info = 0;
	/*tmp1 and tmp2 are used for mid-step sum*/
	double *tmp2=(double *)malloc(num_ensemble*sizeof(double) );
	tmp2 = initial_vector(num_ensemble);

	double *ftmp1 = (double *)malloc(num_ensemble*num_ensemble*sizeof(double));
			ftmp1 = initial_vector(num_ensemble*num_ensemble);
	double *fPa_tilde=(double *)malloc(num_ensemble*num_ensemble*sizeof(double));
	double *fwa_bar = (double *)malloc(num_ensemble*sizeof(double) );
	double *fWa = (double *)malloc(num_ensemble*num_ensemble*sizeof(double) );
	double *finvR=(double *)malloc(radius*radius*sizeof(double) );
	double *C = (double *)malloc(num_ensemble*num_ensemble*sizeof(double) );
			C = initial_vector( num_ensemble*num_ensemble);
	double *tempz = (double *)malloc(num_ensemble*num_ensemble*sizeof(double) );
	double *eigval=(double *)malloc(num_ensemble*sizeof(double));
	double *eigvec=(double *)malloc(num_ensemble*num_ensemble*sizeof(double));
	double *fvecxai=(double *)malloc((radius+1)*num_ensemble*sizeof(double) );
	double *fvecYb = (double *)malloc(num_ensemble*radius*sizeof(double) );
	double *fvecXb = (double *)malloc(num_ensemble*(radius+1)*sizeof(double) );
	double tmp = (num_ensemble-1)/rho;

	double *localR = (double *)malloc(radius*radius*sizeof(double) );
	double *localxb_bar = (double *)malloc((radius+1)*sizeof(double) );
	double *localyb_bar = (double *)malloc(radius*sizeof(double) );
	double *localobs = (double *)malloc(radius*sizeof(double) );
	double **localXb=mydmatrix(num_ensemble,(radius+1));
	double **localYb=mydmatrix(num_ensemble,radius);	
	
	float *dist_pardump = (float *)malloc(numpar*(numpar-1)/2*sizeof(float) );
	float *dist_cdump = (float *)malloc(numline*numpar*sizeof(float) );
	float *dist_datem = (float *)malloc(numobs*numpar*sizeof(float) );
	
	all_dist(dist_pardump, dist_cdump, dist_datem, xb_tlat, xb_tlon, xb_zlvl, yb_clat, yb_clon, obs_lat, obs_lon, numpar, numline, numobs);

	for (i=0; i<numpar; i++)
	{
	if (xb_zlvl[i] < 101.0)
	{
	// GET (update point) Center point in vector for update
	center = find_neighbors(i, radius, num_ensemble, numpar, numline, numobs, xb_bar, yb_bar, en_mass, en_con, obs, localxb_bar, localyb_bar, localXb, localYb, localobs, localR, len_r, dist_pardump, dist_cdump, dist_datem, valid_dist);						
	len_L = *len_r;

	if (len_L != 0)
	{
	// inverse of R
	info = matrix_invert(len_L, localR, finvR);

	mat2fvec(num_ensemble, radius, localYb, fvecYb);
	
	cblas_dgemm(CblasColMajor, CblasNoTrans, CblasNoTrans, num_ensemble, 
				len_L, len_L, 1.0, fvecYb, num_ensemble, finvR, len_L,
				0.0, C, num_ensemble);		
	cblas_dgemm(CblasColMajor, CblasNoTrans, CblasTrans, num_ensemble, 
				num_ensemble,len_L, 1.0, C, num_ensemble, 
				fvecYb, num_ensemble, 0.0, ftmp1, num_ensemble);
	// yb_bar - y_obs
				
	cblas_daxpy(len_L, -1.0, localyb_bar, 1, localobs, 1);
	cblas_dgemv(CblasColMajor, CblasNoTrans, num_ensemble, len_L, 1.0, 
				C, num_ensemble, localobs, 1, 0.0, tmp2, 1);
	for(k=0; k<num_ensemble; k++)
	    ftmp1[k*num_ensemble+k]=ftmp1[k*num_ensemble+k]+tmp;
	// End of computing tmp1=MStep5+tmp1
	//  wa_bar = Pa_tilde*(C*(y_obs-yb_bar))
	info = matrix_sqrt_invert_once(num_ensemble, ftmp1, fPa_tilde, fWa, 
									C, tempz, eigval, eigvec, sqrtcon);
			
	matrix_vector_mult(num_ensemble, num_ensemble, 1.0, 0.0, fPa_tilde, 
						tmp2, fwa_bar);

	for (k=0; k<num_ensemble; k++)
	{//  wai = Wa + repmat(wa_bar,1, num_ensemble);
	cblas_daxpy(num_ensemble, 1.0, fwa_bar, 1, fWa+k*num_ensemble, 1);
	}

	//  Now, wai is stored in matrix Wa
	//  xai = Xb*wai + repmat(xb_bar, 1, num_ensemble)	
	mat2fvec(num_ensemble, radius+1, localXb, fvecXb);	

	cblas_dgemm(CblasColMajor, CblasTrans, CblasNoTrans, radius+1, 
				num_ensemble, num_ensemble, 1.0, fvecXb, 
				num_ensemble, fWa, num_ensemble, 0.0, fvecxai, radius+1);

	for (k=0; k<num_ensemble; k++)
		{
		printf("en_bar[%d][%d] was %12.3E\n",k,i,en_bar[k][i]);	
		en_bar[k][i] = fvecxai[center+k*(radius+1)]
						+ localxb_bar[center];
		printf("en_bar[%d][%d] is %12.3E\n",k,i,en_bar[k][i]);				
	//H_en is stored by row-wise, num_ensemble by lenxy
		}
		} // End of l_len != 0
	else // if l_len == 0
		{
		// Not change the value.
		} //End of l_len == 0
	  }	
	} // End for i

    // End of doing analysis
free(tmp2);free(ftmp1);free(fPa_tilde);
free(dist_pardump); free(dist_cdump); free(dist_datem);
free(fwa_bar); free(fWa); free(C); free(tempz); 
free(eigval); free(eigvec); 
free(fvecxai); free(finvR); 
free(fvecXb);
free(fvecYb); free(localxb_bar);
free(localyb_bar); free(localR);
free(localobs);
free_mydmatrix(localXb);
free_mydmatrix(localYb);

}

int matrix_sqrt_invert_once(int n, double *X, double *Y, double *Z, double *C, 
		double *tempz, double *eigval, double *eigvec, double sqrtcon) 
{
	//double *eigval, *eigvec; //, *temp, *tempz;
	int info = 0, i, j;
	double tmp = 0.0;	
	/*  Memory allocation  */
	
	/*  Eigen decomposition  */
	info=eigen_decomposition(n, X, eigvec, eigval);
	if (info != 0) return info;

	/*  Check for positive definitiveness*/
	for (i=0; i<n; i++) if (eigval[i]<0) {
		fprintf(stderr, "In matrix_square_root_n: Matrix is not PD.\n");
		return 1;
	}
	/*  Build square rooted diagonal matrix, and inverse  */
	for (i=0; i<n; i++) for (j=0; j<n; j++) {Y[i*n+j] = 0.0; Z[i*n+j] = 0.0;}
	for (i=0; i<n; i++) {tmp = 1.0/eigval[i]; Y[i*n+i] = tmp; Z[i*n+i] = sqrt(tmp);}
	
	
	/*  Multiply the eigenvectors with this diagonal matrix Y. Store back in Y  */
	matrix_matrix_mult(n, n, n, 1.0, 0, eigvec, Y, C);
	matrix_transpose(n, n, eigvec, tempz);
	matrix_matrix_mult(n, n, n, 1.0, 0, C, tempz, Y);
	
	matrix_matrix_mult(n, n, n, 1.0, 0, eigvec, Z, C);
	matrix_matrix_mult(n, n, n, 1.0, 0, C, tempz, Z);
	vector_scale(n*n, sqrtcon, Z);		
	/*  Cleanup and exit  */
	//free(eigvec); free(eigval);// free(temp); free(tempz);
	return info;
}
